import pandas as pd
import numpy as np
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import (classification_report, confusion_matrix, 
                            accuracy_score, precision_score, recall_score, 
                            f1_score, roc_auc_score, roc_curve)
from sklearn.model_selection import cross_val_score, GridSearchCV
import matplotlib.pyplot as plt
import seaborn as sns
import warnings
from imblearn.ensemble import BalancedRandomForestClassifier
warnings.filterwarnings('ignore')

# 设置中文字体支持
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

def load_processed_data():
    """加载预处理后的数据"""
    print("=== 加载数据 ===")
    
    try:
        # 加载训练数据
        train_data = pd.read_excel('data_train_balanced.xlsx')
        test_data = pd.read_excel('data_test_validation.xlsx')
        
        print(f"训练集: {len(train_data)}行, {len(train_data.columns)}列")
        print(f"测试集: {len(test_data)}行, {len(test_data.columns)}列")
        
        # 分离特征和标签
        feature_columns = [col for col in train_data.columns if col != '是否异常']
        
        X_train = train_data[feature_columns]
        y_train = train_data['是否异常']
        X_test = test_data[feature_columns]
        y_test = test_data['是否异常']
        
        # 数据基本信息
        print(f"特征数量: {len(feature_columns)}")
        print(f"训练集分布 - 正常: {(y_train == 0).sum()}, 异常: {(y_train == 1).sum()}")
        print(f"测试集分布 - 正常: {(y_test == 0).sum()}, 异常: {(y_test == 1).sum()}")
        
        return X_train, X_test, y_train, y_test, feature_columns
        
    except Exception as e:
        print(f"数据加载失败: {e}")
        return None, None, None, None, None

def train_random_forest(X_train, y_train, random_state=42):
    """训练随机森林模型"""
    print("\n=== 训练随机森林模型 ===")
    
    # 基础随机森林模型
    rf_basic = BalancedRandomForestClassifier(
        n_estimators=500,  # 树的数量，建议比默认多
        max_depth=25,  # 或者设一个上限，比如 25
        sampling_strategy="auto",  # 自动平衡少数类
        replacement=False,  # 下采样时不放回
        random_state=42,
        n_jobs=-1
    )
    
    # 训练模型
    rf_basic.fit(X_train, y_train)
    print("基础随机森林模型训练完成")
    
    # 超参数调优
    print("开始超参数调优...")
    param_grid = {
        'n_estimators': [50, 100],
        'max_depth': [10, 20],
        'min_samples_split': [2, 5],
        'min_samples_leaf': [1, 2]
    }
    
    rf_grid = RandomForestClassifier(random_state=random_state, n_jobs=-1)
    
    # 网格搜索
    grid_search = GridSearchCV(
        rf_grid, param_grid, 
        cv=5, scoring='f1', 
        n_jobs=-1, verbose=1
    )
    
    grid_search.fit(X_train, y_train)
    
    print(f"最佳参数: {grid_search.best_params_}")
    print(f"最佳交叉验证F1得分: {grid_search.best_score_:.4f}")
    
    return rf_basic, grid_search.best_estimator_

def evaluate_model(model, X_train, X_test, y_train, y_test, model_name="随机森林"):
    """评估模型性能"""
    print(f"\n=== {model_name}模型评估 ===")
    
    # 训练集预测
    y_train_pred = model.predict(X_train)
    y_train_proba = model.predict_proba(X_train)[:, 1]
    
    # 测试集预测
    y_test_pred = model.predict(X_test)
    y_test_proba = model.predict_proba(X_test)[:, 1]
    
    # 计算各种指标
    results = {
        '数据集': ['训练集', '测试集'],
        '准确率': [
            accuracy_score(y_train, y_train_pred),
            accuracy_score(y_test, y_test_pred)
        ],
        '精确率': [
            precision_score(y_train, y_train_pred),
            precision_score(y_test, y_test_pred)
        ],
        '召回率': [
            recall_score(y_train, y_train_pred),
            recall_score(y_test, y_test_pred)
        ],
        'F1得分': [
            f1_score(y_train, y_train_pred),
            f1_score(y_test, y_test_pred)
        ],
        'AUC': [
            roc_auc_score(y_train, y_train_proba),
            roc_auc_score(y_test, y_test_proba)
        ]
    }
    
    # 创建结果DataFrame
    results_df = pd.DataFrame(results)
    print("\n性能指标对比:")
    print(results_df.round(4))
    
    return results_df, y_train_pred, y_test_pred, y_train_proba, y_test_proba

def plot_confusion_matrices(y_train, y_test, y_train_pred, y_test_pred):
    """绘制混淆矩阵"""
    print("\n=== 生成混淆矩阵图 ===")
    
    fig, axes = plt.subplots(1, 2, figsize=(12, 5))
    
    # 训练集混淆矩阵
    cm_train = confusion_matrix(y_train, y_train_pred)
    sns.heatmap(cm_train, annot=True, fmt='d', cmap='Blues', ax=axes[0])
    axes[0].set_title('训练集混淆矩阵')
    axes[0].set_xlabel('预测标签')
    axes[0].set_ylabel('真实标签')
    
    # 测试集混淆矩阵
    cm_test = confusion_matrix(y_test, y_test_pred)
    sns.heatmap(cm_test, annot=True, fmt='d', cmap='Blues', ax=axes[1])
    axes[1].set_title('测试集混淆矩阵')
    axes[1].set_xlabel('预测标签')
    axes[1].set_ylabel('真实标签')
    
    plt.tight_layout()
    plt.savefig('confusion_matrices.png', dpi=300, bbox_inches='tight')
    plt.show()
    
    return cm_train, cm_test

def plot_roc_curves(y_train, y_test, y_train_proba, y_test_proba):
    """绘制ROC曲线"""
    print("\n=== 生成ROC曲线图 ===")
    
    # 计算ROC曲线数据
    fpr_train, tpr_train, _ = roc_curve(y_train, y_train_proba)
    fpr_test, tpr_test, _ = roc_curve(y_test, y_test_proba)
    
    auc_train = roc_auc_score(y_train, y_train_proba)
    auc_test = roc_auc_score(y_test, y_test_proba)
    
    plt.figure(figsize=(8, 6))
    plt.plot(fpr_train, tpr_train, label=f'训练集 (AUC = {auc_train:.4f})', linewidth=2)
    plt.plot(fpr_test, tpr_test, label=f'测试集 (AUC = {auc_test:.4f})', linewidth=2)
    plt.plot([0, 1], [0, 1], 'k--', alpha=0.6, label='随机分类')
    
    plt.xlabel('假阳性率 (FPR)')
    plt.ylabel('真阳性率 (TPR)')
    plt.title('ROC曲线对比')
    plt.legend()
    plt.grid(True, alpha=0.3)
    
    plt.savefig('roc_curves.png', dpi=300, bbox_inches='tight')
    plt.show()

def analyze_feature_importance(model, feature_columns, top_n=15):
    """分析特征重要性"""
    print(f"\n=== 特征重要性分析 (Top {top_n}) ===")
    
    # 获取特征重要性
    importances = model.feature_importances_
    
    # 创建特征重要性DataFrame
    feature_importance = pd.DataFrame({
        '特征名': feature_columns,
        '重要性': importances
    }).sort_values('重要性', ascending=False)
    
    print(feature_importance.head(top_n))
    
    # 绘制特征重要性图
    plt.figure(figsize=(10, 8))
    top_features = feature_importance.head(top_n)
    plt.barh(range(top_n), top_features['重要性'], color='skyblue')
    plt.yticks(range(top_n), top_features['特征名'])
    plt.xlabel('特征重要性')
    plt.title(f'Top {top_n} 特征重要性')
    plt.gca().invert_yaxis()
    plt.grid(True, alpha=0.3)
    
    plt.tight_layout()
    plt.savefig('feature_importance.png', dpi=300, bbox_inches='tight')
    plt.show()
    
    return feature_importance

def generate_detailed_report(results_df, cm_train, cm_test, feature_importance):
    """生成详细分析报告"""
    print("\n=== 生成详细分析报告 ===")
    
    report = f"""
# 随机森林分类模型分析报告

## 模型性能总结

### 整体表现
- 训练集准确率: {results_df.loc[0, '准确率']:.4f}
- 测试集准确率: {results_df.loc[1, '准确率']:.4f}
- 训练集F1得分: {results_df.loc[0, 'F1得分']:.4f}
- 测试集F1得分: {results_df.loc[1, 'F1得分']:.4f}
- 训练集AUC: {results_df.loc[0, 'AUC']:.4f}
- 测试集AUC: {results_df.loc[1, 'AUC']:.4f}

### 模型泛化能力分析
- 准确率差异: {abs(results_df.loc[0, '准确率'] - results_df.loc[1, '准确率']):.4f}
- F1得分差异: {abs(results_df.loc[0, 'F1得分'] - results_df.loc[1, 'F1得分']):.4f}
- AUC差异: {abs(results_df.loc[0, 'AUC'] - results_df.loc[1, 'AUC']):.4f}

### 混淆矩阵分析

#### 训练集
- 真阴性: {cm_train[0, 0]}
- 假阳性: {cm_train[0, 1]}
- 假阴性: {cm_train[1, 0]}
- 真阳性: {cm_train[1, 1]}

#### 测试集
- 真阴性: {cm_test[0, 0]}
- 假阳性: {cm_test[0, 1]}
- 假阴性: {cm_test[1, 0]}
- 真阳性: {cm_test[1, 1]}

### Top 10 重要特征
{feature_importance.head(10).to_string(index=False)}

## 结论和建议

### 模型表现
1. 模型在训练集和测试集上的表现{'相对稳定' if abs(results_df.loc[0, 'F1得分'] - results_df.loc[1, 'F1得分']) < 0.1 else '存在一定差异'}
2. {'过拟合风险较低' if abs(results_df.loc[0, 'AUC'] - results_df.loc[1, 'AUC']) < 0.1 else '可能存在过拟合'}

### 实际应用建议
1. 模型可用于孕妇异常检测的辅助诊断
2. 建议结合医学专家判断，不应完全依赖模型结果
3. 可根据特征重要性优化检测流程，重点关注重要特征

生成时间: {pd.Timestamp.now().strftime('%Y-%m-%d %H:%M:%S')}
"""
    
    # 保存报告
    with open('model_analysis_report.txt', 'w', encoding='utf-8') as f:
        f.write(report)
    
    print("详细报告已保存到: model_analysis_report.txt")
    
    return report

def main():
    """主函数"""
    print("=== 随机森林分类模型训练与评估 ===")
    
    # 1. 加载数据
    X_train, X_test, y_train, y_test, feature_columns = load_processed_data()
    if X_train is None:
        return
    
    # 2. 训练模型
    rf_basic, rf_optimized = train_random_forest(X_train, y_train)
    
    # 3. 评估基础模型
    print("\n" + "="*50)
    print("基础随机森林模型评估")
    print("="*50)
    results_basic, y_train_pred_basic, y_test_pred_basic, y_train_proba_basic, y_test_proba_basic = evaluate_model(
        rf_basic, X_train, X_test, y_train, y_test, "基础随机森林"
    )
    
    # 4. 评估优化模型
    print("\n" + "="*50)
    print("优化后随机森林模型评估")
    print("="*50)
    results_optimized, y_train_pred_opt, y_test_pred_opt, y_train_proba_opt, y_test_proba_opt = evaluate_model(
        rf_optimized, X_train, X_test, y_train, y_test, "优化随机森林"
    )
    
    # 5. 选择最佳模型进行详细分析
    best_model = rf_optimized if results_optimized.loc[1, 'F1得分'] > results_basic.loc[1, 'F1得分'] else rf_basic
    best_results = results_optimized if results_optimized.loc[1, 'F1得分'] > results_basic.loc[1, 'F1得分'] else results_basic
    best_name = "优化随机森林" if results_optimized.loc[1, 'F1得分'] > results_basic.loc[1, 'F1得分'] else "基础随机森林"
    
    print(f"\n最佳模型: {best_name}")
    
    # 6. 绘制图表
    if best_name == "优化随机森林":
        cm_train, cm_test = plot_confusion_matrices(y_train, y_test, y_train_pred_opt, y_test_pred_opt)
        plot_roc_curves(y_train, y_test, y_train_proba_opt, y_test_proba_opt)
    else:
        cm_train, cm_test = plot_confusion_matrices(y_train, y_test, y_train_pred_basic, y_test_pred_basic)
        plot_roc_curves(y_train, y_test, y_train_proba_basic, y_test_proba_basic)
    
    # 7. 特征重要性分析
    feature_importance = analyze_feature_importance(best_model, feature_columns)
    
    # 8. 生成详细报告
    report = generate_detailed_report(best_results, cm_train, cm_test, feature_importance)
    
    # 9. 输出总结
    print("\n" + "="*60)
    print("分析完成！生成的文件:")
    print("📊 confusion_matrices.png - 混淆矩阵图")
    print("📈 roc_curves.png - ROC曲线图")
    print("📊 feature_importance.png - 特征重要性图")
    print("📄 model_analysis_report.txt - 详细分析报告")
    print("="*60)

if __name__ == "__main__":
    main()